import numpy as np
import os

def gen_iscomplex_data():
    # 创建输入和输出目录
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    # 定义目标形状：(8, 2048)
    target_shape = (8, 2048)
    total_elements = target_shape[0] * target_shape[1]  # 8*2048=16384个元素

    # 生成复数类型数据（complex64），包含两种情况：
    # 1. 虚部非零（应返回True）
    # 2. 虚部为零（应返回False）
    # 注意：np.iscomplex只关注虚部是否非零，与原始类型无关（只要最终是复数类型）
    
    # 生成实部（统一范围）
    real_part = np.random.uniform(-10, 10, total_elements).astype(np.float32)
    
    # 生成虚部：50%概率非零，50%概率为0
    # 非零虚部使用随机值，零虚部直接置0
    mask_nonzero_imag = np.random.random(total_elements) < 0.5  # 50%概率为True
    imag_part = np.where(
        mask_nonzero_imag,
        np.random.uniform(-5, 5, total_elements).astype(np.float32),  # 非零虚部
        0.0  # 零虚部
    )
    
    # 组合为复数数组（complex64）
    input_flat = real_part + imag_part.astype(np.complex64) * 1j
    input_x = input_flat.reshape(target_shape)  # 重塑为目标形状

    # 生成黄金数据：严格按照np.iscomplex的逻辑（虚部非零则为True）
    golden_flat = np.iscomplex(input_flat)  # 直接使用np.iscomplex生成黄金结果
    golden = golden_flat.reshape(target_shape)

    # 保存为二进制文件
    input_x.tofile("./input/input_x.bin")    # 类型：complex64，每个元素8字节（实部4+虚部4）
    golden.tofile("./output/golden.bin")     # 类型：bool，每个元素1字节（True=1，False=0）


if __name__ == "__main__":
    gen_iscomplex_data()
